use async_ach_spsc as spsc;
use eds_reader::Reader;
use eds_writer::Writer;
use std::collections::BTreeMap;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::UdpSocket;
use unmp_link::prelude::*;

const DATA_LEN: usize = 4096;

struct Conn {
    identifier: String,
    recv_buf: Vec<u8>,
    socket: Arc<UdpSocket>,
    target: SocketAddr,
    data_ch: Arc<spsc::Spsc<Vec<u8>, 8>>,
}
#[async_trait]
impl ConnectionDriver for Conn {
    fn is_subnet(&self) -> bool {
        true
    }
    fn kind(&self) -> &str {
        "udp_server"
    }
    fn id(&self) -> &str {
        &self.identifier
    }
    async fn send(&mut self, buf: Bytes) -> Result<(), SendError> {
        let mut eds_writer = Writer::new(3);
        let mut tx = eds_writer.get_load().unwrap();
        tx.extend_from_slice(&buf);
        let data = eds_writer.get_data(tx);
        match self.socket.send_to(data, self.target).await {
            Ok(_) => Ok(()),
            Err(_) => Err(SendError::Disconnected),
        }
    }
    async fn recv(&mut self) -> Result<Bytes, RecvError> {
        let mut reader = Reader::new(2);
        if !self.recv_buf.is_empty() {
            let l = reader.recv(&self.recv_buf);
            self.recv_buf = Vec::from(&self.recv_buf[l..]);
            if let Some(load) = reader.get_load() {
                return Ok(load);
            }
        }

        loop {
            let data = self.data_ch.take_recver().unwrap().recv().await;
            if data.is_empty() {
                return Err(RecvError::Disconnected);
            }
            let l = reader.recv(&data);
            if let Some(load) = reader.get_load() {
                self.recv_buf = Vec::from(&data[l..]);
                return Ok(load);
            }
        }
    }
}

pub struct Server {
    identifier: String,
    socket: Arc<UdpSocket>,
    data_ch: BTreeMap<SocketAddr, Arc<spsc::Spsc<Vec<u8>, 8>>>,
}
impl Server {
    pub async fn new(addr: SocketAddr) -> Result<Listener, ()> {
        let identifier = addr.to_string();
        let socket = UdpSocket::bind(addr).await.map_err(|_| ())?;
        Ok(Listener::new(Box::new(Self {
            identifier,
            socket: Arc::new(socket),
            data_ch: BTreeMap::new(),
        })))
    }
}
#[async_trait]
impl ListenerDriver for Server {
    fn kind(&self) -> &str {
        "udp_server"
    }
    fn id(&self) -> &str {
        &self.identifier
    }
    async fn accept(&mut self) -> Result<Connection, RecvError> {
        loop {
            let mut buf = Vec::with_capacity(DATA_LEN);
            unsafe { buf.set_len(DATA_LEN) };
            let (len, addr) = match self.socket.recv_from(&mut buf).await {
                Ok(d) => d,
                Err(_) => continue,
            };
            unsafe { buf.set_len(len) };
            if let Some(ch) = self.data_ch.get(&addr) {
                ch.take_sender().unwrap().send(buf).await;
                continue;
            };
            let identifier = addr.to_string();
            let ch = Arc::new(spsc::Spsc::new());
            ch.take_sender().unwrap().send(buf).await;
            self.data_ch.insert(addr, ch.clone());
            return Ok(Connection::new(Box::new(Conn {
                identifier,
                recv_buf: Vec::new(),
                socket: self.socket.clone(),
                target: addr,
                data_ch: ch,
            })));
        }
    }
}
